import os
import pygame
import sys
import numpy as np

# 初始化Pygame
pygame.init()

# 常量定义
WINDOW_SIZE = 800  # 窗口大小
BOARD_SIZE = 15  # 棋盘大小
GRID_SIZE = WINDOW_SIZE // (BOARD_SIZE + 1)  # 格子大小
PIECE_SIZE = int(GRID_SIZE * 0.8)  # 棋子大小

# 颜色定义
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
BROWN = (199, 155, 111)
BUTTON_COLOR = (0, 128, 0)
BUTTON_HOVER_COLOR = (0, 255, 0)
BUTTON_TEXT_COLOR = (255, 255, 255)

# 创建窗口
screen = pygame.display.set_mode((WINDOW_SIZE, WINDOW_SIZE))
pygame.display.set_caption("五子棋")

# 定义玩家和电脑的棋子
player_chess = '★'
computer_chess = 'O'


class GomokuGame:
    def __init__(self):
        self.board = [[' ' for _ in range(BOARD_SIZE)] for _ in range(BOARD_SIZE)]
        self.mark_sheet = [[0 for _ in range(BOARD_SIZE)] for _ in range(BOARD_SIZE)]
        self.current_player = 1  # 1表示玩家，2表示电脑
        self.game_over = False

    def draw_board(self):
        # 填充棋盘背景色
        screen.fill(BROWN)

        # 绘制棋盘网格
        for i in range(BOARD_SIZE):
            # 横线
            pygame.draw.line(screen, BLACK,
                             (GRID_SIZE, (i + 1) * GRID_SIZE),
                             (WINDOW_SIZE - GRID_SIZE, (i + 1) * GRID_SIZE))
            # 竖线
            pygame.draw.line(screen, BLACK,
                             ((i + 1) * GRID_SIZE, GRID_SIZE),
                             ((i + 1) * GRID_SIZE, WINDOW_SIZE - GRID_SIZE))

    def draw_pieces(self):
        for i in range(BOARD_SIZE):
            for j in range(BOARD_SIZE):
                if self.board[i][j] == player_chess:  # 玩家棋子
                    pygame.draw.circle(screen, BLACK,
                                       ((j + 1) * GRID_SIZE, (i + 1) * GRID_SIZE),
                                       PIECE_SIZE // 2)
                elif self.board[i][j] == computer_chess:  # 电脑棋子
                    pygame.draw.circle(screen, WHITE,
                                       ((j + 1) * GRID_SIZE, (i + 1) * GRID_SIZE),
                                       PIECE_SIZE // 2)

    def place_piece(self, row, col, chess_type):
        if self.board[row][col] == ' ' and not self.game_over:
            self.board[row][col] = chess_type
            if self.check_winner(row, col, chess_type):
                self.game_over = True
                if chess_type == player_chess:
                    print("恭喜你赢了！！！")
                else:
                    print("很遗憾你输了！！！")

    def check_winner(self, row, col, chess_type):
        # 检查四个方向：水平、垂直、两个对角线
        directions = [(0, 1), (1, 0), (1, 1), (1, -1)]

        for dx, dy in directions:
            count = 1  # 当前位置已经有一个棋子

            # 正向检查
            x, y = row + dx, col + dy
            while 0 <= x < BOARD_SIZE and 0 <= y < BOARD_SIZE and self.board[x][y] == chess_type:
                count += 1
                x += dx
                y += dy

            # 反向检查
            x, y = row - dx, col - dy
            while 0 <= x < BOARD_SIZE and 0 <= y < BOARD_SIZE and self.board[x][y] == chess_type:
                count += 1
                x -= dx
                y -= dy

            if count >= 5:
                return True
        return False

    def chess_computer(self):
        max_grade = -999
        best_move = None

        # 更新评分表
        self.mark_sheet = [[self.grade(x, y) for y in range(BOARD_SIZE)] for x in range(BOARD_SIZE)]

        for i in range(BOARD_SIZE):
            for j in range(BOARD_SIZE):
                if self.mark_sheet[i][j] > max_grade:
                    max_grade = self.mark_sheet[i][j]
                    best_move = (i, j)

        if best_move:
            self.place_piece(best_move[0], best_move[1], computer_chess)

    def grade(self, x, y):
        # 检查是否是非法落子点
        if x < 4 or x >= BOARD_SIZE - 4 or y < 4 or y >= BOARD_SIZE - 4 or self.board[x][y] != ' ':
            return -999  # 如果是非法点，则扣分

        # 连五子和堵四子优先执行
        for ch in [computer_chess, player_chess]:
            self.board[x][y] = ch
            if self.check_winner(x, y, ch):
                self.board[x][y] = ' '
                return 999
            else:
                self.board[x][y] = ' '

        grade = 0
        flag = 0

        # 检查落子点是否太过于偏远
        # 检查周围一圈是否都没有棋子
        for i in range(-1, 2):
            for j in range(-1, 2):
                if 0 <= x + i < BOARD_SIZE and 0 <= y + j < BOARD_SIZE:
                    if self.board[x + i][y + j] != ' ' and (i != 0 or j != 0):
                        flag = 1
                        break
            if flag:
                break
        if flag == 0:
            grade -= 1
        # 检查周围一圈是否都没有对方棋子
        flag = 0
        for i in range(-1, 2):
            for j in range(-1, 2):
                if 0 <= x + i < BOARD_SIZE and 0 <= y + j < BOARD_SIZE:
                    if self.board[x + i][y + j] == player_chess and (i != 0 or j != 0):
                        flag = 1
                        break
            if flag:
                break
        if flag == 0:
            grade -= 1

        # 检查落子点是否处于地图边界
        if (4 < x <= 6 or BOARD_SIZE - 6 < x <= BOARD_SIZE - 4) and (
                4 < y <= 6 or BOARD_SIZE - 6 < y <= BOARD_SIZE - 4):
            grade -= 3
        if (4 < x <= 6 or BOARD_SIZE - 6 < x <= BOARD_SIZE - 4) or (4 < y <= 6 or BOARD_SIZE - 6 < y <= BOARD_SIZE - 4):
            grade -= 2

        # 电脑的防守与进攻
        flag = 1  # 避免了连四子和连三子的分数累加到一块
        count = 0  # 用于接收可以一次性连成几个的四子，三子，堵住几个的三子

        # 可连四子
        count = self.count4(x, y, computer_chess)
        if count:
            flag = 0
            grade += 8 * count
        # 可连三子
        count = self.count3(x, y, computer_chess)
        if count and flag:
            grade += 4 * count

        flag = 1
        # 可堵四子
        count = self.count4(x, y, player_chess)
        if count:
            flag = 0
            grade += 8 * count
        # 可堵三子
        count = self.count3(x, y, player_chess)
        if count and flag:
            grade += 4 * count

        return grade

    def count4(self, x, y, chess_type):
        count = 0
        size = BOARD_SIZE

        # 检查水平方向
        judge = [0]  # 为了使索引和实际对应，我们让judge在初始时就已经含有一个元素
        if x + 3 < size:
            judge.append(chess_type == self.board[x + 1][y] == self.board[x + 2][y] == self.board[x + 3][y])
        if x - 1 >= 0 and x + 2 < size:
            judge.append(chess_type == self.board[x - 1][y] == self.board[x + 1][y] == self.board[x + 2][y])
        if x - 2 >= 0 and x + 1 < size:
            judge.append(chess_type == self.board[x - 2][y] == self.board[x - 1][y] == self.board[x + 1][y])
        if x - 3 >= 0:
            judge.append(chess_type == self.board[x - 3][y] == self.board[x - 2][y] == self.board[x - 1][y])

        judge_useful = [0]
        if x + 4 < size and x - 1 >= 0:
            judge_useful.append(judge[1] and self.board[x + 4][y] == self.board[x - 1][y] == ' ')
        if x + 3 < size and x - 2 >= 0:
            judge_useful.append(judge[2] and self.board[x - 2][y] == self.board[x + 3][y] == ' ')
        if x + 2 < size and x - 3 >= 0:
            judge_useful.append(judge[3] and self.board[x - 3][y] == self.board[x + 2][y] == ' ')
        if x + 1 < size and x - 4 >= 0:
            judge_useful.append(judge[4] and self.board[x - 4][y] == self.board[x + 1][y] == ' ')

        flag = 0
        if any(judge_useful[1:]):
            count += 2
            flag = 1
        if any(judge[1:]) and flag:
            count += 1

        # 检查竖直方向
        judge = [0]
        if y + 3 < size:
            judge.append(chess_type == self.board[x][y + 1] == self.board[x][y + 2] == self.board[x][y + 3])
        if y - 1 >= 0 and y + 2 < size:
            judge.append(chess_type == self.board[x][y - 1] == self.board[x][y + 1] == self.board[x][y + 2])
        if y - 2 >= 0 and y + 1 < size:
            judge.append(chess_type == self.board[x][y - 2] == self.board[x][y - 1] == self.board[x][y + 1])
        if y - 3 >= 0:
            judge.append(chess_type == self.board[x][y - 3] == self.board[x][y - 2] == self.board[x][y - 1])

        judge_useful = [0]
        if y + 4 < size and y - 1 >= 0:
            judge_useful.append(judge[1] and self.board[x][y + 4] == self.board[x][y - 1] == ' ')
        if y + 3 < size and y - 2 >= 0:
            judge_useful.append(judge[2] and self.board[x][y - 2] == self.board[x][y + 3] == ' ')
        if y + 2 < size and y - 3 >= 0:
            judge_useful.append(judge[3] and self.board[x][y - 3] == self.board[x][y + 2] == ' ')
        if y + 1 < size and y - 4 >= 0:
            judge_useful.append(judge[4] and self.board[x][y - 4] == self.board[x][y + 1] == ' ')

        flag = 0
        if any(judge_useful[1:]):
            count += 2
            flag = 1
        if any(judge[1:]) and flag:
            count += 1

        # 检查正对角线方向
        judge = [0]
        if x + 3 < size and y + 3 < size:
            judge.append(chess_type == self.board[x + 1][y + 1] == self.board[x + 2][y + 2] == self.board[x + 3][y + 3])
        if x + 2 < size and y + 2 < size and x - 1 >= 0 and y - 1 >= 0:
            judge.append(chess_type == self.board[x - 1][y - 1] == self.board[x + 1][y + 1] == self.board[x + 2][y + 2])
        if x + 1 < size and y + 1 < size and x - 2 >= 0 and y - 2 >= 0:
            judge.append(chess_type == self.board[x - 2][y - 2] == self.board[x - 1][y - 1] == self.board[x + 1][y + 1])
        if x - 3 >= 0 and y - 3 >= 0:
            judge.append(chess_type == self.board[x - 3][y - 3] == self.board[x - 2][y - 2] == self.board[x - 1][y - 1])

        judge_useful = [0]
        if x + 4 < size and y + 4 < size and x - 1 >= 0 and y - 1 >= 0:
            judge_useful.append(judge[1] and self.board[x + 4][y + 4] == self.board[x - 1][y - 1] == ' ')
        if x + 3 < size and y + 3 < size and x - 2 >= 0 and y - 2 >= 0:
            judge_useful.append(judge[2] and self.board[x - 2][y - 2] == self.board[x + 3][y + 3] == ' ')
        if x + 2 < size and y + 2 < size and x - 3 >= 0 and y - 3 >= 0:
            judge_useful.append(judge[3] and self.board[x - 3][y - 3] == self.board[x + 2][y + 2] == ' ')
        if x + 1 < size and y + 1 < size and x - 4 >= 0 and y - 4 >= 0:
            judge_useful.append(judge[4] and self.board[x - 4][y - 4] == self.board[x + 1][y + 1] == ' ')

        flag = 0
        if any(judge_useful[1:]):
            count += 2
            flag = 1
        if any(judge[1:]) and flag:
            count += 1

        # 检查反对角线方向
        judge = [0]
        if x - 3 >= 0 and y + 3 < size:
            judge.append(chess_type == self.board[x - 1][y + 1] == self.board[x - 2][y + 2] == self.board[x - 3][y + 3])
        if x + 2 < size and y - 2 >= 0 and x - 1 >= 0 and y + 1 < size:
            judge.append(chess_type == self.board[x + 1][y - 1] == self.board[x - 1][y + 1] == self.board[x - 2][y + 2])
        if x + 1 < size and y - 1 >= 0 and x - 2 >= 0 and y + 2 < size:
            judge.append(chess_type == self.board[x + 2][y - 2] == self.board[x + 1][y - 1] == self.board[x - 1][y + 1])
        if x + 3 < size and y - 3 >= 0:
            judge.append(chess_type == self.board[x + 3][y - 3] == self.board[x + 2][y - 2] == self.board[x + 1][y - 1])

        judge_useful = [0]
        if x - 4 >= 0 and y + 4 < size and x + 1 < size and y - 1 >= 0:
            judge_useful.append(judge[1] and self.board[x - 4][y + 4] == self.board[x + 1][y - 1] == ' ')
        if x - 3 >= 0 and y + 3 < size and x + 2 < size and y - 2 >= 0:
            judge_useful.append(judge[2] and self.board[x + 2][y - 2] == self.board[x - 3][y + 3] == ' ')
        if x - 2 >= 0 and y + 2 < size and x + 3 < size and y - 3 >= 0:
            judge_useful.append(judge[3] and self.board[x + 3][y - 3] == self.board[x - 2][y + 2] == ' ')
        if x - 1 >= 0 and y + 1 < size and x + 4 < size and y - 4 >= 0:
            judge_useful.append(judge[4] and self.board[x + 4][y - 4] == self.board[x - 1][y + 1] == ' ')

        flag = 0
        if any(judge_useful[1:]):
            count += 2
            flag = 1
        if any(judge[1:]) and flag:
            count += 1

        return count

    def count3(self, x, y, chess_type):
        count = 0
        size = BOARD_SIZE

        # 检查水平方向
        judge = [0]
        if x + 2 < size:
            judge.append(chess_type == self.board[x + 1][y] == self.board[x + 2][y])
        if x - 1 >= 0 and x + 1 < size:
            judge.append(chess_type == self.board[x - 1][y] == self.board[x + 1][y])
        if x - 2 >= 0:
            judge.append(chess_type == self.board[x - 2][y] == self.board[x - 1][y])

        judge_useful = [0]
        if x + 3 < size and x - 1 >= 0:
            judge_useful.append(judge[1] and self.board[x + 3][y] == self.board[x - 1][y] == ' ')
        if x + 2 < size and x - 2 >= 0:
            judge_useful.append(judge[2] and self.board[x - 2][y] == self.board[x + 2][y] == ' ')
        if x + 1 < size and x - 3 >= 0:
            judge_useful.append(judge[3] and self.board[x - 3][y] == self.board[x + 1][y] == ' ')

        flag = 0
        if any(judge_useful[1:]):
            count += 2
            flag = 1
        if any(judge[1:]) and flag:
            count += 1

        # 检查竖直方向
        judge = [0]
        if y + 2 < size:
            judge.append(chess_type == self.board[x][y + 1] == self.board[x][y + 2])
        if y - 1 >= 0 and y + 1 < size:
            judge.append(chess_type == self.board[x][y - 1] == self.board[x][y + 1])
        if y - 2 >= 0:
            judge.append(chess_type == self.board[x][y - 2] == self.board[x][y - 1])

        judge_useful = [0]
        if y + 3 < size and y - 1 >= 0:
            judge_useful.append(judge[1] and self.board[x][y + 3] == self.board[x][y - 1] == ' ')
        if y + 2 < size and y - 2 >= 0:
            judge_useful.append(judge[2] and self.board[x][y - 2] == self.board[x][y + 2] == ' ')
        if y + 1 < size and y - 3 >= 0:
            judge_useful.append(judge[3] and self.board[x][y - 3] == self.board[x][y + 1] == ' ')

        flag = 0
        if any(judge_useful[1:]):
            count += 2
            flag = 1
        if any(judge[1:]) and flag:
            count += 1

        # 检查正对角线方向
        judge = [0]
        if x + 2 < size and y + 2 < size:
            judge.append(chess_type == self.board[x + 1][y + 1] == self.board[x + 2][y + 2])
        if x + 1 < size and y + 1 < size and x - 1 >= 0 and y - 1 >= 0:
            judge.append(chess_type == self.board[x - 1][y - 1] == self.board[x + 1][y + 1])
        if x - 2 >= 0 and y - 2 >= 0:
            judge.append(chess_type == self.board[x - 2][y - 2] == self.board[x - 1][y - 1])

        judge_useful = [0]
        if x + 3 < size and y + 3 < size and x - 1 >= 0 and y - 1 >= 0:
            judge_useful.append(judge[1] and self.board[x + 3][y + 3] == self.board[x - 1][y - 1] == ' ')
        if x + 2 < size and y + 2 < size and x - 2 >= 0 and y - 2 >= 0:
            judge_useful.append(judge[2] and self.board[x - 2][y - 2] == self.board[x + 2][y + 2] == ' ')
        if x + 1 < size and y + 1 < size and x - 3 >= 0 and y - 3 >= 0:
            judge_useful.append(judge[3] and self.board[x - 3][y - 3] == self.board[x + 1][y + 1] == ' ')

        flag = 0
        if any(judge_useful[1:]):
            count += 2
            flag = 1
        if any(judge[1:]) and flag:
            count += 1

        # 检查反对角线方向
        judge = [0]
        if x - 2 >= 0 and y + 2 < size:
            judge.append(chess_type == self.board[x - 1][y + 1] == self.board[x - 2][y + 2])
        if x + 1 < size and y - 1 >= 0 and x - 1 >= 0 and y + 1 < size:
            judge.append(chess_type == self.board[x + 1][y - 1] == self.board[x - 1][y + 1])
        if x + 2 < size and y - 2 >= 0:
            judge.append(chess_type == self.board[x + 2][y - 2] == self.board[x + 1][y - 1])

        judge_useful = [0]
        if x - 3 >= 0 and y + 3 < size and x + 1 < size and y - 1 >= 0:
            judge_useful.append(judge[1] and self.board[x - 3][y + 3] == self.board[x + 1][y - 1] == ' ')
        if x - 2 >= 0 and y + 2 < size and x + 2 < size and y - 2 >= 0:
            judge_useful.append(judge[2] and self.board[x + 2][y - 2] == self.board[x - 2][y + 2] == ' ')
        if x - 1 >= 0 and y + 1 < size and x + 3 < size and y - 3 >= 0:
            judge_useful.append(judge[3] and self.board[x + 3][y - 3] == self.board[x - 1][y + 1] == ' ')

        flag = 0
        if any(judge_useful[1:]):
            count += 2
            flag = 1
        if any(judge[1:]) and flag:
            count += 1

        return count


def draw_button(text, x, y, width, height, action=None):
    mouse = pygame.mouse.get_pos()
    click = pygame.mouse.get_pressed()

    if x + width > mouse[0] > x and y + height > mouse[1] > y:
        pygame.draw.rect(screen, BUTTON_HOVER_COLOR, (x, y, width, height))
        if click[0] == 1 and action is not None:
            action()
    else:
        pygame.draw.rect(screen, BUTTON_COLOR, (x, y, width, height))

    font = pygame.font.SysFont(None, 40)
    text_surf = font.render(text, True, BUTTON_TEXT_COLOR)
    text_rect = text_surf.get_rect(center=(x + width / 2, y + height / 2))
    screen.blit(text_surf, text_rect)


def game_intro():
    intro = True

    while intro:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()

        screen.fill(BROWN)
        font = pygame.font.SysFont(None, 100)
        text_surf = font.render("五子棋", True, BUTTON_TEXT_COLOR)
        text_rect = text_surf.get_rect(center=(WINDOW_SIZE / 2, WINDOW_SIZE / 4))
        screen.blit(text_surf, text_rect)

        draw_button("玩家先手", WINDOW_SIZE / 4, WINDOW_SIZE / 2, 200, 50, player_first)
        draw_button("电脑先手", WINDOW_SIZE / 2, WINDOW_SIZE / 2, 200, 50, computer_first)

        pygame.display.flip()
        pygame.time.Clock().tick(15)


def player_first():
    global game
    game = GomokuGame()
    game_loop()


def computer_first():
    global game
    game = GomokuGame()
    game.chess_computer()
    game_loop()


def game_loop():
    clock = pygame.time.Clock()

    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()

            if event.type == pygame.MOUSEBUTTONDOWN and not game.game_over:
                x, y = event.pos
                col = round((x - GRID_SIZE) / GRID_SIZE)
                row = round((y - GRID_SIZE) / GRID_SIZE)
                if 0 <= row < BOARD_SIZE and 0 <= col < BOARD_SIZE:
                    game.place_piece(row, col, player_chess)
                    if not game.game_over:
                        game.chess_computer()

        game.draw_board()
        game.draw_pieces()
        pygame.display.flip()
        clock.tick(60)


if __name__ == "__main__":
    game_intro()